import random
from server.inference import qwen2_5_7b_inference, qwen2_5_14b_inference
import requests
import json

def qwen_7b_model_reply(info, prompt):
    url = "http://127.0.0.1:8001/task/chat_qwen_7b"
    payload = {
        "info": info,
        "prompt": prompt
    }
    # 发送POST请求
    response = requests.post(url, json=payload)
    print("千问7b的回复", response)
    return response

def qwen_14b_model_reply(info, prompt):
    url = "http://127.0.0.1:8002/task/chat_qwen_14b"
    payload = {
        "info": info,
        "prompt": prompt
    }
    # 发送POST请求
    response = requests.post(url, json=payload)
    print("千问14b的回复", response)
    return response

# agent
import json

def agent_init(info, prompt):
    state = 'True'
    chat_history = []
    max_attempts = 3  # 限制最多尝试3次
    attempts = 0  # 初始化计数器

    while state == 'True' and attempts < max_attempts:
        # 获取两个模型的回复
        qwen_7b_response = qwen_7b_model_reply(info, prompt)
        qwen_14b_response = qwen_14b_model_reply(info, prompt)

        # 构建14B模型的输入，包含两个回复
        new_prompt = (
            f"提问: {prompt}\n"
            f"答案1: {qwen_7b_response}\n"
            f"答案2: {qwen_14b_response}\n"
            f"请从答案1或答案2中选择一个更好的答案。'answer' 字段是你在上述答案1和答案2中选择的内容，'state' 字段如果认为答案可以输出，请填入 True；如果认为不可以输出，请填入 False，按照如下的 JSON 格式返回：\n"
            f'{{\n\t"answer": "your_answer_here",\n\t"state": "your_state_here"\n}}'
        )

        # 打印当前的提示模板
        print("agent模板: ", new_prompt)

        # 追加当前的new_prompt到聊天历史中
        chat_history.append(new_prompt)

        # 将历史记录加回新提示中，并作为14B模型的新输入
        combined_prompt = "\n".join(chat_history)  # 将所有聊天历史组合成一个输入
        print("combined_prompt: ", combined_prompt)
        
        # 调用14B模型的回复
        response = qwen_14b_model_reply(info, combined_prompt)

        # 尝试解析14B模型返回的JSON格式数据
        try:
            response_data = json.loads(response)
            state = response_data.get("state", "False")  # 更新state状态
            final_answer = response_data.get("answer", "")  # 获取最终的回复
        except json.JSONDecodeError:
            print("无法解析14B模型的回复为JSON，直接返回14b的回复: ", response)
            return qwen_14b_response

        # 如果14B模型返回True，则输出最终答案并结束循环
        if state == "True":
            print("模型认为可以输出的最终答案: ", final_answer)
            return final_answer

        # 更新尝试次数
        attempts += 1

    # 如果超过3次，强制输出最后一个模型的答案
    print(f"已达到最大尝试次数 {max_attempts} 次，强制输出最终答案。")
    return final_answer




